package com.heiye.blog.ai.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @author: heiye
 * @date: 2025/10/09 下午9:50
 * @version: v1.0.0
 * @description: 自定义线程池
 */
@Configuration
public class ThreadPoolConfig {
    /**
     * HTTP 请求线程池（IO 密集型任务）
     *
     * @return
     */
    @Bean("httpRequestExecutor")
    public ThreadPoolTaskExecutor httpRequestExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        // 核心线程数（保持常驻）
        executor.setCorePoolSize(50);
        // 最大线程数（突发流量时扩容）
        executor.setMaxPoolSize(200);
        // 任务队列容量（缓冲突发请求）
        executor.setQueueCapacity(1000);
        // 空闲线程存活时间（秒）
        executor.setKeepAliveSeconds(120);
        // 线程名前缀（便于监控）
        executor.setThreadNamePrefix("http-fetcher-");
        // 拒绝策略（由调用线程执行）
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        // 初始化线程池
        executor.initialize();
        return executor;
    }

    /**
     * 结果处理线程池（CPU 密集型任务）
     *
     * @return
     */
    @Bean("resultProcessingExecutor")
    public ThreadPoolTaskExecutor resultProcessingExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        // 核心线程数（等于CPU核心数）
        executor.setCorePoolSize(Runtime.getRuntime().availableProcessors());
        // 最大线程数（不超过CPU核心数2倍）
        executor.setMaxPoolSize(Runtime.getRuntime().availableProcessors() * 2);
        // 较小队列（避免任务堆积）
        executor.setQueueCapacity(200);
        // 线程名前缀（便于监控）
        executor.setThreadNamePrefix("result-processor-");
        // 拒绝策略（由调用线程执行）
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        // 初始化线程池
        executor.initialize();
        return executor;
    }

    /**
     * AI 任务线程池
     *
     * @return
     */
    @Bean(name = "aiTaskExecutor")
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        // 核心线程数
        executor.setCorePoolSize(10);
        // 最大线程数
        executor.setMaxPoolSize(50);
        // 队列容量
        executor.setQueueCapacity(200);
        // 线程活跃时间（秒）
        executor.setKeepAliveSeconds(30);
        // 线程名前缀
        executor.setThreadNamePrefix("AiExecutor-");

        // 拒绝策略：由调用线程处理（一般为主线程）
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

        // 等待所有任务结束后再关闭线程池
        executor.setWaitForTasksToCompleteOnShutdown(true);
        // 设置等待时间，如果超过这个时间还没有销毁就强制销毁，以确保应用最后能够被关闭，而不是被没有完成的任务阻塞
        executor.setAwaitTerminationSeconds(60);

        executor.initialize();
        return executor;
    }
}
